
package instrumentation;
import java.util.Map.Entry;
import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.logging.Logger;

import Operations.ValuesGenerator;


import random.TestClusterBuilder;
import soot.ActiveBodyTransformer;
import soot.Body;
import soot.BodyTransformer;
import soot.Local;
import soot.PackManager;
import soot.PatchingChain;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Transform;
import soot.Trap;
import soot.Unit;
import soot.jimple.GotoStmt;
import soot.jimple.IdentityStmt;
import soot.jimple.IfStmt;
import soot.jimple.Jimple;
import soot.jimple.JimpleBody;
import soot.jimple.LookupSwitchStmt;
import soot.jimple.Stmt;
import soot.jimple.TableSwitchStmt;
import soot.jimple.toolkits.scalar.NopEliminator;
import soot.tagkit.LineNumberTag;
import soot.tagkit.StringTag;
import soot.util.Chain;
import soot.SootMain;
import instrumentation.WhiteInstrumenter;
import util.GAConfig;

public class Instrumenter {

	public static interface UnifiedInstrumentator {

		public int init(Body oldBody, Body newBody, Chain<Unit> newUnits, IdentityStmt[] paramDefs);

		public void processPre(Chain<Unit> newUnits, Stmt op);

		public int processPost(Chain<Unit> newUnits, Stmt op);
		public void done();
	}

	private static final Logger logger = Logger.getLogger("GAMD.instrumenter");


	private static final boolean instrumenter        = true;
	private static final boolean nopEliminator       = true;

	public static void prepare( Set<Class<?>>  classToInstrument) {
		String[] SOOT_CONF = new String[] { "-validate", "-keep-line-number", "-f", "c", "-output-dir", GAConfig.getDirInstrumented().toString() };

//		logger.info("Instrumenting: " + toInstrument);

		String params[] = new String[SOOT_CONF.length + 2 + classToInstrument.size()];

		params[0] = "-cp";
		
		params[1] =GAConfig.getDirCompiled().toString();// "C:/Experiment/Instrumenter/./bin";

		params[1] += File.pathSeparator + System.getProperty("java.class.path");
		params[1] += File.pathSeparator + System.getProperty("sun.boot.class.path");

		int i = 2;
		for(String s : SOOT_CONF)
			params[i++] = s;
		for(Class<?> clz : classToInstrument)
			params[i++] = clz.getName();

		logger.config("Launching SOOT with command line parameters:\n" + Arrays.toString(params));
		System.out.println("\nLaunching SOOT with command line parameters:\n" + Arrays.toString(params));
		SootMain.singleton.processCmdLine(params);
	}

	@SuppressWarnings({ "unused", "unchecked" })
	public static void run(   Set<Class<?>>  toInstrument) {

		ClassesInstrumenter instr = null;

		if(instrumenter) {
		instr = new ClassesInstrumenter();

			Set<String> done = new HashSet<String>();
		//	while(!toInstrument.isEmpty()) {
			for(Class<?> clz:toInstrument){
				String className = clz.getName();

				Scene.v().loadClassAndSupport(className);
				SootClass sClass = Scene.v().getSootClass(className);
				System.out.println();
			}
		}

		String last = null;

		if(instrumenter) {
			String newPhase = "jtp.coverageInstrumenter";
			logger.fine("Enabled phase: " + newPhase);
			if(last == null) PackManager.v().getPack("jtp").add(new Transform(newPhase, instr));
			last = newPhase;
		}
		if(nopEliminator) {
			String newPhase = "jtp.nopEliminator";
			logger.fine("Enabled phase: " + newPhase);
			if(last == null) PackManager.v().getPack("jtp").add(new Transform(newPhase, ActiveBodyTransformer.v(NopEliminator.v())));
			else PackManager.v().getPack("jtp").insertAfter(new Transform(newPhase, ActiveBodyTransformer.v(NopEliminator.v())), last);
			last = newPhase;
		}



		SootMain.singleton.run();

	//	if(instrumenter)
			wi.done();
	}

	private static class ClassesInstrumenter extends BodyTransformer {


		/** java.lang.Throwable soot class */
		public static final SootClass throwableClass;

		static {
			Scene.v().loadClassAndSupport(Throwable.class.getName());
			throwableClass = Scene.v().getSootClass(Throwable.class.getName());
		}
		public ClassesInstrumenter() {
		//	this.wi=wi;
		}

		/**
		 * Preprocess a class. In this phase it is allowed to add methods to the class.
		 * Notice that new methods will be processed in subsequent phases. To skip them (or to
		 * skip pre-existent methods) annotate them with testful.utils.Skip.
		 * @param sClass the class being preprocessed
		 */

		@Override
		@SuppressWarnings("rawtypes")
		protected void internalTransform(Body oldBody, String phaseName, Map options) {
			final SootMethod method = oldBody.getMethod();
			final String methodName = method.getName();
	//		if(methodName.equals("<init>")||
			if(methodName.equals("<clinit>") ) return ;

			final SootClass sClass = method.getDeclaringClass();
			final Iterator<Unit> oldStmtIt = oldBody.getUnits().snapshotIterator();

			final JimpleBody newBody = Jimple.v().newBody(method);
			method.setActiveBody(newBody);
			final PatchingChain<Unit> newUnits = newBody.getUnits();
			newBody.getLocals().addAll(oldBody.getLocals());

			final Local exc = Jimple.v().newLocal("__throwable_exc__", throwableClass.getType());
			newBody.getLocals().add(exc);

			System.out.println("\n\n** Instrumenting " + sClass.getName() + "." + methodName + "("+ method.getParameterTypes()+ ")");

			/** stores the start of an operation (i.e. nopPre) */
			final Map<Unit, Unit> start = new HashMap<Unit, Unit>();

			// instrumentation structure:
			//   initial method code (@this=this, params, superCall)
			//
			// :NOP_PRE_INIT
			//   initial tracker code
			// :NOP_POST_INIT
			//
			// :NOP_BEGIN ( try { )
			//
			// ---- the operation is NOT an identity statement ----
			// :NOP_PRE
			// :NOP_PRE_TRACK
			//   tracking code pre
			// :NOP_ORIG
			//   original code
			// :NOP_POST_TRACK
			//   tracking code post
			// :NOP_AFTER
			//
			// ---- the operation IS an identity statement ----
			// :NOP_PRE
			// :NOP_ORIG
			//   original code
			// :NOP_PRE_TRACK
			//   tracking code
			// :NOP_POST_TRACK
			//   tracking code
			// :NOP_AFTER
			//
			// :NOP_END ( } catch(Throwable exc) { )
			//   exceptional tracker code
			//   throw exc

			// --------------------------------------------------------------------------
			// initial method code (@this=this, params, superCall)
			// --------------------------------------------------------------------------

			// skip special statements: this
			if(!method.isStatic()) newUnits.add(oldStmtIt.next());

			// skip special statements: params
			int nParams = method.getParameterCount();
			final IdentityStmt[] paramDefs = new IdentityStmt[nParams];
			for(int i = 0; i < nParams; i++) {
				Unit paramDef = oldStmtIt.next();
				newUnits.add(paramDef);
				paramDefs[i] = (IdentityStmt) paramDef;
			}
			// --------------------------------------------------------------------------
			// initialization
			// --------------------------------------------------------------------------

			{
				final Unit nopPre = Jimple.v().newNopStmt();
				nopPre.addTag(new StringTag("nopInitPre"));
				newUnits.add(nopPre);

	//			for(UnifiedInstrumentator i : instrumenters)
				startBranch=wi.init(oldBody, newBody, newUnits, paramDefs);

				final Unit nopPost = Jimple.v().newNopStmt();
				nopPost.addTag(new StringTag("nopInitAfter"));
				newUnits.add(nopPost);
			}

			final Unit nopMethodBegin = Jimple.v().newNopStmt();
			nopMethodBegin.addTag(new StringTag("nopMethodPre"));
			newUnits.add(nopMethodBegin);

			// --------------------------------------------------------------------------
			// statement instrumentation
			// --------------------------------------------------------------------------

			while(oldStmtIt.hasNext()) {
				Stmt stmt = (Stmt) oldStmtIt.next();
				System.out.println("\tstmt "+ stmt.toString());
				final Unit nopPre = Jimple.v().newNopStmt();
				nopPre.addTag(new StringTag("nopPre"));
				start.put(stmt, nopPre);
				newUnits.add(nopPre);

				final Unit nopPreTrack = Jimple.v().newNopStmt();
				nopPreTrack.addTag(new StringTag("nopPreTrack"));

				final Unit nopOrig = Jimple.v().newNopStmt();
				nopOrig.addTag(new StringTag("nopOrig"));

				final Unit nopPostTrack = Jimple.v().newNopStmt();
				nopPostTrack.addTag(new StringTag("nopPostTrack"));

				final Unit nopAfter = Jimple.v().newNopStmt();
				nopAfter.addTag(new StringTag("nopAfter"));

				final Stmt newStmt = (Stmt) stmt.clone();
				LineNumberTag line = (LineNumberTag) stmt.getTag("LineNumberTag");
				if(line != null) 
					newStmt.addTag(line);
				if(stmt instanceof IdentityStmt) {

					// insert original stmt
					newUnits.add(nopOrig);
					newUnits.add(newStmt);

					// preprocess
					newUnits.add(nopPreTrack);
//					for(UnifiedInstrumentator i : instrumenters)
					wi.processPre(newUnits, stmt);

					// postprocess
					newUnits.add(nopPostTrack);
//					for(UnifiedInstrumentator i : instrumenters)
						wi.processPost(newUnits, stmt);

				} else {

					// preprocess
					newUnits.add(nopPreTrack);
//					for(UnifiedInstrumentator i : instrumenters)
					wi.processPre(newUnits, stmt);

					// insert original stmt
					newUnits.add(nopOrig);
					newUnits.add(newStmt);

					// postprocess
					newUnits.add(nopPostTrack);
//					for(UnifiedInstrumentator i : instrumenters)
					endBranch=wi.processPost(newUnits, stmt);
				}

				newUnits.add(nopAfter);
			}
			if(startBranch!=endBranch){
//				calls.add(new String("Java"));
				MethodBranchs mBranch=new MethodBranchs(methodName,startBranch,endBranch-1,calls);
				methodBranchInfo.put(methodName+method.getParameterTypes(), mBranch);
			}
			calls.clear();
			final Unit nopMethodEnd = Jimple.v().newNopStmt();
			nopMethodEnd.addTag(new StringTag("nopMethodEnd"));
			newUnits.add(nopMethodEnd);

			// Fix jumps (goto, if, switch)
			for(Unit unit : newUnits) {
				if(unit instanceof GotoStmt) {
					GotoStmt gotoStmt = (GotoStmt) unit;
					Unit newTarget = start.get(gotoStmt.getTarget());
	//				System.out.println("GotoStmt: "+newTarget);
					if(newTarget != null) gotoStmt.setTarget(newTarget);
				} else if(unit instanceof IfStmt) {
					IfStmt ifStmt = (IfStmt) unit;
					Unit newTarget = start.get(ifStmt.getTarget());
//					System.out.println("IfStmt: "+newTarget);
					if(newTarget != null) ifStmt.setTarget(newTarget);
				} else if(unit instanceof TableSwitchStmt) {
					TableSwitchStmt sw = (TableSwitchStmt) unit;

					Unit newTarget = start.get(sw.getDefaultTarget());
					if(newTarget != null) sw.setDefaultTarget(newTarget);

					final int lowIndex = sw.getLowIndex();
					for(int idx = 0; idx <= sw.getHighIndex()-lowIndex; idx++) {
						newTarget = start.get(sw.getTarget(idx));
						if(newTarget != null) sw.setTarget(idx, newTarget);
					}
				} else if(unit instanceof LookupSwitchStmt) {
					LookupSwitchStmt sw = (LookupSwitchStmt) unit;

					Unit newTarget = start.get(sw.getDefaultTarget());
					if(newTarget != null) sw.setDefaultTarget(newTarget);

					for(int i = 0; i < sw.getTargetCount(); i++) {
						newTarget = start.get(sw.getTarget(i));
						if(newTarget != null) sw.setTarget(i, newTarget);
					}
				}
			}

			// fix traps (try-catch)
			for(Trap trap : oldBody.getTraps()) {
				final Unit newBegin = start.get(trap.getBeginUnit());
				final Unit newEnd = start.get(trap.getEndUnit());
				final Unit newHandler = start.get(trap.getHandlerUnit());

				newBody.getTraps().add(Jimple.v().newTrap(trap.getException(), newBegin, newEnd, newHandler));
			}


			// --------------------------------------------------------------------------
			// Exceptional tracking activities
			// --------------------------------------------------------------------------
//			final Unit nopCatchBegin = Jimple.v().newNopStmt();
//			nopCatchBegin.addTag(new StringTag("nopCatchPre"));
//			newUnits.add(nopCatchBegin);
//			newUnits.add(Jimple.v().newIdentityStmt(exc, Jimple.v().newCaughtExceptionRef()));
//			// update the traps
//			newBody.getTraps().add(Jimple.v().newTrap(throwableClass, nopMethodBegin, nopMethodEnd, nopCatchBegin));
////			for(UnifiedInstrumentator i : instrumenters)
//				wi.exceptional(newUnits, exc);
//			// The last istruction of the handler is "throw exc"
//			newUnits.add(Jimple.v().newThrowStmt(exc));
//			final Unit nopPost = Jimple.v().newNopStmt();
//			nopPost.addTag(new StringTag("nopCatchAfter"));
//			newUnits.add(nopPost);
		}
	}
	 public static final String lineSep = System.getProperty("line.separator");
	public static void write(File baseDir,String name) {
		try {
			FileWriter fstream = new FileWriter(baseDir+"/"+ name.replace('.', File.separatorChar) + ".txt");

			  BufferedWriter out = new BufferedWriter(fstream);
			  StringBuilder sb = new StringBuilder();
			    for(Entry<String, MethodBranchs> entry : methodBranchInfo.entrySet()) {
			    	sb = new StringBuilder();
			        sb.append(entry.getKey());
			        MethodBranchs mBranch=(MethodBranchs) entry.getValue();
			        sb.append(mBranch.toString());
					 sb.append(lineSep); 
					 out.write(sb.toString());
			    }
			    sb = new StringBuilder();
		        sb.append("#".toString());
				 sb.append(lineSep); 
				 out.write(sb.toString());

			    for (Entry<String, MethodBranchs> entry : methodBranchInfo
					.entrySet()) {
				sb = new StringBuilder();
				sb.append(entry.getKey());
				MethodBranchs mBranch = (MethodBranchs) entry.getValue();
				if (mBranch.getCalls().size() != 0) {
					sb.append(mBranch.callsToString());
					sb.append(lineSep);
					out.write(sb.toString());
				}
			}

			  out.close();
			  
		} catch(IOException e) {
			System.out.println(  e);
		} 
	}

	public static void writeToTxtFile(Set<Class<?>> classesToInstrument){
		for(Class<?> clz:classesToInstrument)
			write(GAConfig.getDirInstrumented(),clz.getName());
	}
	public static  Map<String,MethodBranchs> methodBranchInfo= new LinkedHashMap<String,MethodBranchs>();
	public static int startBranch,endBranch;
	public static List<String> calls= new ArrayList<String>();;
	 public  static   WhiteInstrumenter wi=new WhiteInstrumenter();  
	public static void main(String[] args) 
	    {
		try {
			ValuesGenerator.resetSeed(10);
			 Set<Class<?>> allClasses = new LinkedHashSet<Class<?>>();

		      allClasses.addAll(TestClusterBuilder.readClassFromFile());

		      System.out.println("There are " + allClasses.size() + " classes as input");
		      
		      System.out.println(allClasses);
		      

		    Set<Class<?>> classesToInstrument = TestClusterBuilder.filterUntestableClasses(allClasses);
	//		List<String> toInstrument=new ArrayList<String>(classesToTest); 
		    prepare(classesToInstrument);
		    run(classesToInstrument);

		    writeToTxtFile(classesToInstrument);
		    System.out.println("Done instrumented!");
		} catch (Exception e) {
			System.err.println("Error during the instrumentation: " + e.getMessage());
		//	System.exit(1);
		}	   
	    }
}
